{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SFTTrainerを使用した教師あり微調整\n",
    "\n",
    "このノートブックでは、`trl`ライブラリの`SFTTrainer`を使用して`HuggingFaceTB/SmolLM2-135M`モデルを微調整する方法を示します。ノートブックのセルを実行すると、モデルの微調整が行われます。さまざまなデータセットを試して、難易度を選択できます。\n",
    "\n",
    "<div style='background-color: lightblue; padding: 10px; border-radius: 5px; margin-bottom: 20px; color:black'>\n",
    "    <h2 style='margin: 0;color:blue'>演習: SFTTrainerを使用したSmolLM2の微調整</h2>\n",
    "    <p>Hugging Faceのリポジトリからデータセットを取得し、それを使用してモデルを微調整します。</p>\n",
    "    <p><b>難易度レベル</b></p>\n",
    "    <p>🐢 `HuggingFaceTB/smoltalk`データセットを使用</p>\n",
    "    <p>🐕 `bigcode/the-stack-smol`データセットを試し、特定のサブセット`data/python`でコード生成モデルを微調整します。</p>\n",
    "    <p>🦁 実際の使用ケースに関連するデータセットを選択します。</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Google Colabでの要件のインストール\n",
    "# !pip install transformers datasets trl huggingface_hub\n",
    "\n",
    "\n",
    "# Hugging Faceへの認証\n",
    "from huggingface_hub import login\n",
    "\n",
    "login()\n",
    "\n",
    "# 便利のため、Hugging Faceのトークンを.envファイルのHF_TOKENとして環境変数に設定できます"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 必要なライブラリをインポート\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "from datasets import load_dataset\n",
    "from trl import SFTConfig, SFTTrainer, setup_chat_format\n",
    "import torch\n",
    "\n",
    "# デバイスを動的に設定\n",
    "device = (\n",
    "    \"cuda\"\n",
    "    if torch.cuda.is_available()\n",
    "    else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n",
    ")\n",
    "\n",
    "# モデルとトークナイザーをロード\n",
    "model_name = \"HuggingFaceTB/SmolLM2-135M\"\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    pretrained_model_name_or_path=model_name\n",
    ").to(device)\n",
    "tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=model_name)\n",
    "\n",
    "# チャット形式を設定\n",
    "model, tokenizer = setup_chat_format(model=model, tokenizer=tokenizer)\n",
    "\n",
    "# 微調整されたモデルを保存および/またはアップロードするための名前を設定\n",
    "finetune_name = \"SmolLM2-FT-MyDataset\"\n",
    "finetune_tags = [\"smol-course\", \"module_1\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ベースモデルでの生成\n",
    "\n",
    "ここでは、チャットテンプレートを持たないベースモデルを試してみます。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# トレーニング前のベースモデルを試してみましょう\n",
    "prompt = \"プログラミングについての俳句を書いてください\"\n",
    "\n",
    "# テンプレートでフォーマット\n",
    "messages = [{\"role\": \"user\", \"content\": prompt}]\n",
    "formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False)\n",
    "\n",
    "# 応答を生成\n",
    "inputs = tokenizer(formatted_prompt, return_tensors=\"pt\").to(device)\n",
    "outputs = model.generate(**inputs, max_new_tokens=100)\n",
    "print(\"トレーニング前:\")\n",
    "print(tokenizer.decode(outputs[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## データセットの準備\n",
    "\n",
    "例としてデータセットをロードし、トレーニング用にフォーマットします。データセットは、各入力がプロンプトで、出力がモデルの期待される応答である入力-出力ペアで構成されている必要があります。\n",
    "\n",
    "**TRLは、モデルのチャットテンプレートに従って入力メッセージをフォーマットします。** これらは、`role`と`content`キーを持つ辞書のリストとして表現される必要があります。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 例としてデータセットをロード\n",
    "from datasets import load_dataset\n",
    "\n",
    "# TODO: パスと名前のパラメータを使用してデータセットと設定を定義\n",
    "ds = load_dataset(path=\"HuggingFaceTB/smoltalk\", name=\"everyday-conversations\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: 🦁 データセットがTRLでチャットテンプレートに変換できない形式の場合、処理が必要です。 [モジュール](../chat_templates.md)を参照してください。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SFTTrainerの設定\n",
    "\n",
    "`SFTTrainer`は、トレーニングプロセスを制御するさまざまなパラメータで設定されます。これには、トレーニングステップの数、バッチサイズ、学習率、および評価戦略が含まれます。これらのパラメータを特定の要件と計算リソースに応じて調整します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SFTTrainerを設定\n",
    "sft_config = SFTConfig(\n",
    "    output_dir=\"./sft_output\",  # トレーニング結果の出力ディレクトリ\n",
    "    max_steps=1000,  # データセットのサイズとトレーニングの所要時間に応じて調整\n",
    "    per_device_train_batch_size=4,  # GPUのメモリ容量に応じて調整\n",
    "    learning_rate=5e-5,  # 微調整の一般的な開始点\n",
    "    logging_steps=10,  # トレーニングメトリックのログ記録頻度\n",
    "    save_steps=100,  # モデルのチェックポイント保存頻度\n",
    "    evaluation_strategy=\"steps\",  # 定期的な間隔でモデルを評��\n",
    "    eval_steps=50,  # 評価の頻度\n",
    "    use_mps_device=(\n",
    "        True if device == \"mps\" else False\n",
    "    ),  # 混合精度トレーニングにMPSを使用\n",
    "    hub_model_id=finetune_name,  # モデルに一意の名前を付ける\n",
    ")\n",
    "\n",
    "# SFTTrainerを初期化\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    args=sft_config,\n",
    "    train_dataset=ds[\"train\"],\n",
    "    eval_dataset=ds[\"test\"],\n",
    ")\n",
    "\n",
    "# TODO: 🦁 🐕 選択したデータセットに合わせてSFTTrainerのパラメータを調整します。例えば、`bigcode/the-stack-smol`データセットを使用する場合、`content`列を選択する必要があります。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## モデルのトレーニング\n",
    "\n",
    "トレーナーが設定されたので、モデルのトレーニングを進めることができます。トレーニングプロセスは、データセットを反復し、損失を計算し、この損失を最小化するためにモデルのパラメータを更新することを含みます。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# モデルをトレーニング\n",
    "trainer.train()\n",
    "\n",
    "# モデルを保存\n",
    "trainer.save_model(f\"./{finetune_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.push_to_hub(tags=finetune_tags)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style='background-color: lightblue; padding: 10px; border-radius: 5px; margin-bottom: 20px; color:black'>\n",
    "    <h2 style='margin: 0;color:blue'>追加の演習: 微調整されたモデルでの生成</h2>\n",
    "    <p>🐕 微調整されたモデルを使用して、ベースモデルの例と同じように応答を生成します。</p>\n",
    "</div>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 同じメッセージで微調整されたモデルをテスト\n",
    "\n",
    "# トレーニング前のベースモデルを試してみましょう\n",
    "prompt = \"プログラミングについての俳句を書いてください\"\n",
    "\n",
    "# テンプレートでフォーマット\n",
    "messages = [{\"role\": \"user\", \"content\": prompt}]\n",
    "formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False)\n",
    "\n",
    "# 応答を生成\n",
    "inputs = tokenizer(formatted_prompt, return_tensors=\"pt\").to(device)\n",
    "\n",
    "# TODO: 微調整されたモデルを使用して、ベースモデルの例と同じように応答を生成します。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 💐 完了しました！\n",
    "\n",
    "このノートブックでは、`SFTTrainer`を使用して`HuggingFaceTB/SmolLM2-135M`モデルを微調整するためのステップバイステップガイドを提供しました。これらの手順に従うことで、特定のタスクをより効果的に実行するためにモデルを適応させることができます。このコースを続けて作業したい場合、以下の提案を試してみてください：\n",
    "\n",
    "- より高い難易度レベルでこのノートブックを試してみてください。\n",
    "- 同僚のPRをレビューしてください。\n",
    "- IssueやPRを通じてコースの資料を改善してください。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py310",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
